package org.example.brave.dubbo.provider.filter;

import brave.Span;
import brave.Tracer;
import brave.Tracing;
import brave.propagation.CurrentTraceContext;
import brave.propagation.Propagation;
import brave.propagation.TraceContext;
import brave.sampler.SamplerFunction;
import org.apache.dubbo.common.extension.Activate;
import org.apache.dubbo.rpc.*;
import zipkin2.reporter.AsyncReporter;
import zipkin2.reporter.brave.ZipkinSpanHandler;
import zipkin2.reporter.okhttp3.OkHttpSender;

import java.util.Map;

@Activate
public class BraveTraceFilter implements Filter {
    private Tracing tracing;
    private SamplerFunction<Invocation> sampler;

    private void init() {
        if (tracing != null) {
            return;
        }
        sampler = new SamplerFunction<Invocation>() {
            @Override
            public Boolean trySample(Invocation invocation) {
                return !invocation.getMethodName().equalsIgnoreCase("getMetadataInfo");
            }
        };
        tracing = Tracing.newBuilder()
                .localServiceName("brave-dubbo-provider")
                .addSpanHandler(
                        ZipkinSpanHandler.create(
                                AsyncReporter.builder(
                                        OkHttpSender.newBuilder()
                                                .endpoint("http://localhost:9411/api/v2/spans")
                                                .build())
                                        .build()))
                .build();
    }

    Propagation.RemoteGetter<Map<String, Object>> GETTER = new Propagation.RemoteGetter<Map<String, Object>>() {
        @Override
        public Span.Kind spanKind() {
            return Span.Kind.SERVER;
        }

        @Override
        public String get(Map<String, Object> serverAttachment, String key) {
            return (String) serverAttachment.get(key);
        }
    };

    Propagation.RemoteSetter<Map<String, Object>> SETTER = new Propagation.RemoteSetter<Map<String, Object>>() {
        @Override
        public Span.Kind spanKind() {
            return Span.Kind.CLIENT;
        }

        @Override
        public void put(Map<String, Object> clientAttachment, String key, String value) {
            clientAttachment.put(key, value);
        }
    };

    @Override
    public Result invoke(Invoker<?> invoker, Invocation invocation) throws RpcException {
        init();
        Result result = null;
        if (isConsumer(invocation)) {
            TraceContext.Injector<Map<String, Object>> injector = tracing.propagation().injector(SETTER);
            CurrentTraceContext currentTraceContext = tracing.currentTraceContext();
            Tracer tracer = tracing.tracer();
            Span span = tracer.nextSpan().name(invocation.getMethodName()).start();
            try (CurrentTraceContext.Scope ws = currentTraceContext.newScope(span.context())) {
                injector.inject(span.context(), RpcContext.getClientAttachment().getObjectAttachments());
                span.kind(Span.Kind.CLIENT);
                System.out.println("consumer filter test...");
                result = invoker.invoke(invocation);
            } catch (Exception e) {
                span.error(e);
                e.printStackTrace();
            } finally {
                span.finish();
            }
        } else {
            CurrentTraceContext currentTraceContext = tracing.currentTraceContext();
            // ServerAttachment接收到客户端传递过来的参数
            Map<String, Object> serverAattachment = RpcContext.getServerAttachment().getObjectAttachments();
            TraceContext context = (TraceContext) serverAattachment.get(TraceContext.class.getName());
            if (context != null) {
                CurrentTraceContext.Scope scope = currentTraceContext.maybeScope(context);
                try {
                    result = invoker.invoke(invocation);
                } catch (Exception e) {
                    e.printStackTrace();
                } finally {
                    scope.close();
                }
                return result;
            }

            TraceContext.Extractor<Map<String, Object>> extractor = tracing.propagation().extractor(GETTER);
            Tracer tracer = tracing.tracer();
            if (!sampler.trySample(invocation)) {
                return invoker.invoke(invocation);
            }
            Span span = tracer
                    .nextSpan(extractor.extract(RpcContext.getServerAttachment().getObjectAttachments()))
                    .name(invocation.getMethodName());
            // 1. Start the span and add trace headers to the request
            span.start();
            // 2. Put the span in scope so things like log integration works
            CurrentTraceContext.Scope scope = currentTraceContext.newScope(span.context());
            try  {
                result = invoker.invoke(invocation); // 3. Invoke the request
            } catch (Exception e) {
                span.error(e); // 4. Catch any errors
                e.printStackTrace();
            } finally {
                // 5. Complete the span
                span.kind(Span.Kind.SERVER);
                span.finish();
                scope.close();
            }
        }
        return result;
    }

    private boolean isConsumer(Invocation invocation) {
        Invoker<?> invoker = invocation.getInvoker();
        return invoker.getUrl().getParameter("side", "provider").equals("consumer");
    }
}
